
"""
【节点名称】：
    SingleObjectTrackerNode
【依赖项安装】：
    pip install spirems
    git clone -b spirecv2-dds https://gitee.com/amovlab/SpireCV.git
    cd SpireCV && pip install -e .
【订阅类型】：
    sensor_msgs::CompressedImage （输入图像）
    sensor_msgs::CameraCalibration (输入相机参数)
【发布类型】：
    spirecv_msgs::2DTargets （检测结果）
【构造参数说明】：
    parameter_file (str): 全局参数文件
    sms_shutdown (bool): 是否接收全局关闭信号,如果需要长期后台执行,建议设置为False
    specified_input_topic (str): 指定输入的话题地址
    specified_output_topic (str): 指定输出的话题地址
【节点参数】：
    algorithm (str): 单目标跟踪算法, 如"mil"、"goturn"、"DaSiamRPN"、"NanoTrack"
【备注】：
    单目标跟踪算法所需模型需提前下载，然后默认加载模型文件地址为当前文件夹
    For usage download models by following links
    For GOTURN:
        goturn.prototxt and goturn.caffemodel: https://github.com/opencv/opencv_extra/tree/c4219d5eb3105ed8e634278fad312a1a8d2c182d/testdata/tracking
    For DaSiamRPN:
        network:     https://www.dropbox.com/s/rr1lk9355vzolqv/dasiamrpn_model.onnx?dl=0
        kernel_r1:   https://www.dropbox.com/s/999cqx5zrfi7w4p/dasiamrpn_kernel_r1.onnx?dl=0
        kernel_cls1: https://www.dropbox.com/s/qvmtszx5h339a0w/dasiamrpn_kernel_cls1.onnx?dl=0
    For NanoTrack:
        nanotrack_backbone: https://github.com/HonglinChu/SiamTrackers/blob/master/NanoTrack/models/nanotrackv2/nanotrack_backbone_sim.onnx
        nanotrack_headneck: https://github.com/HonglinChu/SiamTrackers/blob/master/NanoTrack/models/nanotrackv2/nanotrack_head_sim.onnx 
"""


# Python 2/3 compatibility
#from __future__ import print_function
#!/usr/bin/env python
# -*- coding: utf-8 -*-
import sys
import os
import yaml
import math
import threading
import cv2 as cv
from queue import Queue
import numpy as np
from typing import Union
from spirems import Subscriber, Publisher, cvimg2sms, sms2cvimg, def_msg, QoS
from spirecv.base.BaseNode import BaseNode
from spirecv.dataloader.CameraNode import CameraNode
from spirecv.dataloader.VideoFileNode import VideoFileNode
from spirecv.visualizer.DetectionVisNode import DetectionVisNode



class SingleObjectTrackerNode(threading.Thread, BaseNode):
    def __init__(
        self,
        job_name: str,
        ip: str = '127.0.0.1',
        port: int = 9094,
        param_dict_or_file: Union[dict, str] = None,
        sms_shutdown: bool = True,
        specified_input_topic: str = None,
        specified_output_topic: str = None,
        realtime_det: bool = True
    ):
        threading.Thread.__init__(self)
        BaseNode.__init__(self, self.__class__.__name__, job_name, ip=ip, port=port,param_dict_or_file=param_dict_or_file, sms_shutdown=sms_shutdown)
        input_url = '/' + job_name + '/sensor/image_raw'
        if specified_input_topic is not None:
            input_url = specified_input_topic
            
        input_url_c = '/' + job_name + '/sensor/calibration_info'
        if specified_output_topic is not None:
            input_url_c = specified_output_topic

        output_url = '/' + job_name + '/detector/results'
        if specified_output_topic is not None:
            output_url = specified_output_topic

        self.realtime_det = realtime_det
        self.job_queue = Queue()
        self.queue_pool.append(self.job_queue)
        
        self.b_clicked = True
        self.b_renew_roi = False
        self.b_begin_track = False
        self.pt_origin = (0, 0)
        self.rect_sel  = (0, 0, 0, 0)
        self.trackerAlgorithm = self.get_param("algorithm", 'nanotrack')
        self.tracker = self.createTracker()
        

        self._image_reader = Subscriber(
            input_url, 'sensor_msgs::CompressedImage', self.image_callback,
            ip=ip, port=port, qos=QoS.Reliability
        )
        self._calibration_reader = Subscriber(
            input_url_c, 'sensor_msgs::CameraCalibration', self.calibration_callback,
            ip=ip, port=port, qos=QoS.Reliability
        )
        
        # self._show_writer = Publisher(
        #    '/' + job_name + '/detector/image_results', 'sensor_msgs::CompressedImage',
        #    ip=ip, port=port
        # )
        
        self._result_writer = Publisher(
            output_url, 'spirecv_msgs::2DTargets',
            ip=ip, port=port, qos=QoS.Reliability
        )
        
        #self._next_writer = Publisher(
        #    '/' + job_name + '/ ', 'std_msgs::Boolean',
        #   ip=ip, port=port
        #)

        self.start()

    def release(self):
        BaseNode.release(self)
        self._image_reader.kill()
        self._calibration_reader.kill()
        self._result_writer.kill()
        # self._show_writer.kill()
        # self._next_writer.kill()
        self._image_reader.join()
        self._calibration_reader.join()
        self._result_writer.join()
        # self._show_writer.join()
        # self._next_writer.join()

    def image_callback(self, msg):
        if self.realtime_det:
            while not self.job_queue.empty():
                self.job_queue.get()
        self.job_queue.put(msg)

    def calibration_callback(self, msg):
        self.width = msg['width']
        self.height = msg['height']
        
        self.camera_matrix = np.array(msg['K']).reshape(3,3)
        self.distortion = np.array(msg['D'])

    def createTracker(self):
        if self.trackerAlgorithm == 'mil':
            tracker = cv.TrackerMIL_create()
        elif self.trackerAlgorithm == 'goturn':
            params = cv.TrackerGOTURN_Params()
            params.modelTxt = "goturn.prototxt" 
            params.modelBin = "goturn.caffemodel"
            tracker = cv.TrackerGOTURN_create(params)
        elif self.trackerAlgorithm == 'dasiamrpn':
            params = cv.TrackerDaSiamRPN_Params()
            params.model = "dasiamrpn_model.onnx"
            params.kernel_cls1 = "dasiamrpn_kernel_cls1.onnx"
            params.kernel_r1 = "dasiamrpn_kernel_r1.onnx"
            tracker = cv.TrackerDaSiamRPN_create(params)
        elif self.trackerAlgorithm == 'nanotrack':
            params = cv.TrackerNano_Params()
            params.backbone = "nanotrack_backbone_sim.onnx"
            params.neckhead = "nanotrack_head_sim.onnx"
            tracker = cv.TrackerNano_create(params)
        else:
            sys.exit("Tracker {} is not recognized. Please use one of three available: mil, goturn, dasiamrpn, nanotrack.".format(self.trackerAlgorithm))
        return tracker
    
    def on_mouse(self, event, x, y, flags, param):
        # param 参数在这里没有被使用，但它是 setMouseCallback 的一部分
        # 我们可以通过闭包或全局变量来传递 self 实例，但这里我们使用类属性
        if event == cv.EVENT_LBUTTONDOWN:
            self.b_clicked = True
            self.pt_origin = (x, y)
            self.rect_sel = (x, y, 0, 0)  # 使用元组表示矩形
        elif event == cv.EVENT_MOUSEMOVE and self.b_clicked:
            # 更新矩形区域
            self.rect_sel = (min(self.pt_origin[0], x), min(self.pt_origin[1], y),
                             abs(x - self.pt_origin[0]), abs(y - self.pt_origin[1]))
        elif event == cv.EVENT_LBUTTONUP:
            self.b_clicked = False
            # 检查矩形大小
            if self.rect_sel[2] * self.rect_sel[3] >= 64:
                # 可以在这里添加处理选中区域的代码
                print("Selected ROI:", self.rect_sel)
                self.b_renew_roi = True

    def run(self):
        
        cv.namedWindow("img")
        cv.setMouseCallback("img", self.on_mouse)
        
        while self.is_running():
            msg = self.job_queue.get(block=True)
            if msg is None:
                break
            
            img = sms2cvimg(msg)
            img = cv.resize(img, (self.width, self.height))
            
            res_msg = def_msg('spirecv_msgs::2DTargets')
            
            res_msg["height"] = self.width
            res_msg["width"] = self.height

            if self.b_renew_roi:
                try:
                    self.tracker.init(img, self.rect_sel)
                    self.b_renew_roi = False
                    self.b_begin_track = True
                except Exception as e:
                    print('Unable to initialize tracker with requested bounding box. Is there any object?')
                    print(e)
                    print('Try again ...')
            
            if self.b_begin_track:
                ok, newbox = self.tracker.update(img)
                res_msg["targets"] = []
                if ok:
                    ann = dict()
                    ann["category_name"] = "TID_" + str(0)
                    ann["bbox"] = newbox
                    res_msg["targets"].append(ann)
                    cv.rectangle(img, newbox, (200,0,0))
       
            cv.imshow('img', img)
            k = cv.waitKey(1)
            if k == ord('c') : 
                self.b_renew_roi = False
                self.b_begin_track = False
    
            self._result_writer.publish(res_msg)
            # END

        self.release()
        print('{} quit!'.format(self.__class__.__name__))

if __name__ == '__main__':

    job_name = 'detection'
    parameter_file = '/home/amov/Desktop/SpireCV/params/spirecv2/default_params.json'
    dataloader = CameraNode(job_name, param_dict_or_file=parameter_file)
    algorithm = SingleObjectTrackerNode(job_name, param_dict_or_file=parameter_file)
